'''
@author: zhangkai
@license: (C) Copyright 2017-2023
@contact: jeffcobile@gmail.com
@Software : PyCharm
@file: danet.py
@time: 2020-06-08 13:51:26
@desc: 
'''
from ELib.model.model_zoo import MODEL_ZOO
from ELib.model.segbase import SegBaseModel
from ELib.basic_model import _FCNHead, PAM_Module, CAM_Module
import torch


@MODEL_ZOO.register()
class DANet(SegBaseModel):
    r"""DANet model from the paper `"Dual Attention Network for Scene Segmentation"
    <https://arxiv.org/abs/1809.02983.pdf>`
    """

    def __init__(self, cfg):
        super(DANet, self).__init__(cfg)
        self.head = DANetHead(2048, self.nclass)
        if self.aux:
            self.auxlayer = _FCNHead(728, self.nclass)
        self.__setattr__('decoder', ['head', 'auxlayer'] if self.aux else ['head'])

    def forward(self, x):
        imsize = x.size()[2:]
        _, _, c3, c4 = self.encoder(x)

        x = self.head(c4)
        x = list(x)
        x[0] = torch.nn.functional.interpolate(x[0], imsize, mode='bilinear', align_corners=True)
        x[1] = torch.nn.functional.interpolate(x[1], imsize, mode='bilinear', align_corners=True)
        x[2] = torch.nn.functional.interpolate(x[2], imsize, mode='bilinear', align_corners=True)

        outputs = list()
        outputs.append(x[0])
        outputs.append(x[1])
        outputs.append(x[2])

        return tuple(outputs)


class DANetHead(torch.nn.Module):
    def __init__(self, in_channels, out_channels, norm_layer=torch.nn.BatchNorm2d):
        super(DANetHead, self).__init__()
        inter_channels = in_channels // 4
        self.conv5a = torch.nn.Sequential(torch.nn.Conv2d(in_channels, inter_channels, 3, padding=1, bias=False),
                                    norm_layer(inter_channels),
                                    torch.nn.ReLU())

        self.conv5c = torch.nn.Sequential(torch.nn.Conv2d(in_channels, inter_channels, 3, padding=1, bias=False),
                                    norm_layer(inter_channels),
                                    torch.nn.ReLU())

        self.sa = PAM_Module(inter_channels)
        self.sc = CAM_Module(inter_channels)
        self.conv51 = torch.nn.Sequential(torch.nn.Conv2d(inter_channels, inter_channels, 3, padding=1, bias=False),
                                    norm_layer(inter_channels),
                                    torch.nn.ReLU())
        self.conv52 = torch.nn.Sequential(torch.nn.Conv2d(inter_channels, inter_channels, 3, padding=1, bias=False),
                                    norm_layer(inter_channels),
                                    torch.nn.ReLU())

        self.conv6 = torch.nn.Sequential(torch.nn.Dropout2d(0.1, False), torch.nn.Conv2d(512, out_channels, 1))
        self.conv7 = torch.nn.Sequential(torch.nn.Dropout2d(0.1, False), torch.nn.Conv2d(512, out_channels, 1))

        self.conv8 = torch.nn.Sequential(torch.nn.Dropout2d(0.1, False), torch.nn.Conv2d(512, out_channels, 1))

    def forward(self, x):
        feat1 = self.conv5a(x)
        sa_feat = self.sa(feat1)
        sa_conv = self.conv51(sa_feat)
        sa_output = self.conv6(sa_conv)

        feat2 = self.conv5c(x)
        sc_feat = self.sc(feat2)
        sc_conv = self.conv52(sc_feat)
        sc_output = self.conv7(sc_conv)

        feat_sum = sa_conv + sc_conv

        sasc_output = self.conv8(feat_sum)

        output = [sasc_output]
        output.append(sa_output)
        output.append(sc_output)
        return tuple(output)